"""教学用途的 VPN 服务器实现。

此模块提供一个通过 TLS 保护的 TCP 转发器，充当教学型 VPN 的“服务器”端。
客户端使用自定义的 JSON 控制报文进行连接，认证成功后，服务器会打开到
请求上游主机与端口的 TCP 连接。随后服务器只需在客户端与上游套接字之间
转发字节，即可在互联网上提供加密传输。

该实现强调可读性与可追踪性，便于学习理解。它**并非**生产环境可用的方案，
缺少真实 VPN 设备应有的诸多特性（高级认证、多路复用、高级路由、UDP 支持等）。
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import ssl
from dataclasses import dataclass
from typing import Optional

# JSON 控制报文的最大字节数，防止恶意或异常客户端在握手完成前发送超大元数据。
MAX_CONTROL_BYTES = 4096


@dataclass
class ServerConfig:
    """VPN 服务器的运行时配置。"""

    bind_host: str
    bind_port: int
    auth_token: str
    certfile: str
    keyfile: str
    cafile: Optional[str]
    require_client_cert: bool


async def relay_stream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    """持续地将 *reader* 中的数据传送到 *writer*，直到遇到 EOF。"""

    # 当源流读取到 EOF 或抛出异常时函数结束。写端的关闭由调用方协调。

    try:
        while True:
            data = await reader.read(65536)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    except asyncio.CancelledError:
        raise
    except Exception as err:  # pragma: no cover - 尽力记录日志
        logging.debug("relay_stream terminated due to %s", err, exc_info=True)


async def handle_client(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    config: ServerConfig,
) -> None:
    """处理新接受的 VPN 客户端连接。"""

    peer = writer.get_extra_info("peername")
    logging.info("Accepted VPN client %s", peer)

    upstream_reader: Optional[asyncio.StreamReader] = None
    upstream_writer: Optional[asyncio.StreamWriter] = None

    try:
        control_bytes = await asyncio.wait_for(reader.readline(), timeout=10)
        if len(control_bytes) > MAX_CONTROL_BYTES:
            raise ValueError("Control message exceeds size limit")

        control_payload = json.loads(control_bytes.decode("utf-8"))
        logging.debug("Received control payload: %s", control_payload)

        token = control_payload.get("token")
        if token != config.auth_token:
            raise PermissionError("Invalid authentication token")

        target_host = control_payload.get("host")
        target_port = control_payload.get("port")

        if not isinstance(target_host, str) or not isinstance(target_port, int):
            raise ValueError("Control message missing 'host' or 'port'")

        # 尝试连接上游套接字。
        logging.info("Connecting to upstream %s:%s for %s", target_host, target_port, peer)
        upstream_reader, upstream_writer = await asyncio.open_connection(target_host, target_port)

        # 通知客户端隧道已准备就绪。
        writer.write(json.dumps({"status": "ok"}).encode("utf-8") + b"\n")
        await writer.drain()

        # 在任一端关闭之前执行双向数据转发。
        await asyncio.gather(
            relay_stream(reader, upstream_writer),
            relay_stream(upstream_reader, writer),
        )
    except (json.JSONDecodeError, UnicodeDecodeError):
        logging.warning("Malformed control message from %s", peer)
    except PermissionError as auth_err:
        logging.warning("Authentication failure from %s: %s", peer, auth_err)
        writer.write(json.dumps({"status": "error", "message": "auth_failed"}).encode("utf-8") + b"\n")
    except (asyncio.TimeoutError, ValueError) as validation_err:
        logging.warning("Invalid control handshake from %s: %s", peer, validation_err)
        writer.write(json.dumps({"status": "error", "message": "bad_request"}).encode("utf-8") + b"\n")
    except ConnectionError as conn_err:
        logging.warning("Failed to reach upstream for %s: %s", peer, conn_err)
        writer.write(json.dumps({"status": "error", "message": "upstream_unreachable"}).encode("utf-8") + b"\n")
    except Exception as err:  # pragma: no cover - 未预期的异常路径
        logging.exception("Unhandled server error for %s", peer)
        writer.write(json.dumps({"status": "error", "message": "internal_error"}).encode("utf-8") + b"\n")
    finally:
        # 尝试优雅地关闭双方套接字。
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:  # pragma: no cover - 尽力完成清理
            logging.debug("Error while closing client writer", exc_info=True)

        if upstream_writer is not None:
            try:
                upstream_writer.close()
                await upstream_writer.wait_closed()
            except Exception:  # pragma: no cover - 记录关闭上游写端的异常
                logging.debug("Error while closing upstream writer", exc_info=True)

        logging.info("Closed session for %s", peer)


def build_ssl_context(config: ServerConfig) -> ssl.SSLContext:
    """配置服务器监听使用的 TLS 上下文。"""

    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    context.load_cert_chain(certfile=config.certfile, keyfile=config.keyfile)

    if config.cafile:
        context.load_verify_locations(config.cafile)

    if config.require_client_cert:
        context.verify_mode = ssl.CERT_REQUIRED
    else:
        context.verify_mode = ssl.CERT_OPTIONAL

    # 默认禁用弱加密套件，在保持配置简单的同时为教学环境增加一点加固。
    context.set_ciphers("ECDHE+AESGCM:ECDHE+CHACHA20")
    return context


async def main_async(config: ServerConfig) -> None:
    ssl_ctx = build_ssl_context(config)

    server = await asyncio.start_server(
        lambda r, w: handle_client(r, w, config),
        host=config.bind_host,
        port=config.bind_port,
        ssl=ssl_ctx,
    )

    addresses = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
    logging.info("VPN server listening on %s", addresses)

    async with server:
        await server.serve_forever()


def parse_args() -> ServerConfig:
    parser = argparse.ArgumentParser(description="教学用途的 TLS VPN 服务器")
    parser.add_argument("--bind-host", default="0.0.0.0", help="服务器绑定的网络接口")
    parser.add_argument("--bind-port", type=int, default=8443, help="服务器监听的端口")
    parser.add_argument("--auth-token", required=True, help="客户端共享的认证令牌")
    parser.add_argument("--cert", required=True, dest="certfile", help="服务器证书文件路径 (PEM)")
    parser.add_argument("--key", required=True, dest="keyfile", help="服务器私钥文件路径 (PEM)")
    parser.add_argument("--ca-cert", dest="cafile", help="用于验证客户端证书的 CA 证书包")
    parser.add_argument(
        "--require-client-cert",
        action="store_true",
        help="强制客户端提供由 --ca-cert 信任的 TLS 证书",
    )
    parser.add_argument("--log-level", default="INFO", help="Python 日志等级（默认：INFO）")

    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s | %(levelname)s | %(message)s",
    )

    return ServerConfig(
        bind_host=args.bind_host,
        bind_port=args.bind_port,
        auth_token=args.auth_token,
        certfile=args.certfile,
        keyfile=args.keyfile,
        cafile=args.cafile,
        require_client_cert=args.require_client_cert,
    )


def main() -> None:
    config = parse_args()

    try:
        asyncio.run(main_async(config))
    except KeyboardInterrupt:
        logging.info("Received shutdown signal. Bye!")


if __name__ == "__main__":
    main()
